import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt

class FeedforwardNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(FeedforwardNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.LeakyReLU(negative_slope=0.01)  # 使用 LeakyReLU
        self.fc2 = nn.Linear(hidden_size, hidden_size)  # 增加一个隐藏层
        self.fc3 = nn.Linear(hidden_size, output_size)
        self.dropout = nn.Dropout(p=0.5)  # 加入 Dropout 防止过拟合

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)  # Dropout
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x

# 重新初始化网络和优化器
input_size = 10
hidden_size = 64  # 更大的隐藏层
output_size = 2

model = FeedforwardNN(input_size, hidden_size, output_size)

# 使用 Adam 优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 使用 CrossEntropyLoss 作为损失函数
criterion = nn.CrossEntropyLoss()

# 创建数据
X_train = torch.randn(100, input_size)
y_train = torch.randint(0, 2, (100,))

train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

# 训练模型
num_epochs = 100
for epoch in range(num_epochs):
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 测试模型
X_test = X_train[:10]  # 选择前10个训练数据样本进行预测
y_test = y_train[:10]  # 选择前10个训练数据标签作为真实标签
model.eval()
with torch.no_grad():
    predictions = model(X_test)
    predicted_labels = torch.argmax(predictions, dim=1)

# 可视化结果
plt.figure(figsize=(10, 6))
plt.plot(y_test.numpy(), label='True Labels', marker='o', linestyle='--', color='b')
plt.plot(predicted_labels.numpy(), label='Predicted Labels', marker='x', linestyle='-', color='r')
plt.xlabel('Sample Index')
plt.ylabel('Label')
plt.title('True vs Predicted Labels')
plt.legend()
plt.show()
